{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Inference Demo\n",
    "This is a demo for loading YTVIS data and performing inference. We have selected multiple trajectories from the YTVIS dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from mmengine.config import Config, ConfigDict\n",
    "from dataset.youtube_loader import VISDataset, vis_collate_fn\n",
    "from mmtrack.registry import DATASETS\n",
    "from mmengine import build_from_cfg\n",
    "from PIL import Image, ImageDraw, ImageSequence, ImageFont\n",
    "from mmtrack.datasets import VideoSampler, EntireVideoBatchSampler\n",
    "\n",
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "classes_vis = ['airplane', 'bear', 'bird', 'boat', 'car',\n",
    "                'cat', 'cow', 'deer', 'dog', 'duck',\n",
    "                'earless_seal', 'elephant', 'fish',\n",
    "                'flying_disc', 'fox', 'frog', 'giant_panda',\n",
    "                'giraffe', 'horse', 'leopard', 'lizard',\n",
    "                'monkey', 'motorbike', 'mouse', 'parrot',\n",
    "                'person', 'rabbit', 'shark', 'skateboard',\n",
    "                'snake', 'snowboard', 'squirrel', 'surfboard',\n",
    "                'tennis_racket', 'tiger', 'train', 'truck',\n",
    "                'turtle', 'whale', 'zebra']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load dataset\n",
    "cfg = Config.fromfile('dataset/youtube_cfg_480.py')\n",
    "val_dataset = build_from_cfg(cfg.val_dataloader.dataset, DATASETS)\n",
    "sampler = VideoSampler(val_dataset)\n",
    "batch_sampler = EntireVideoBatchSampler(sampler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "all_test_case = []\n",
    "def get_fixed_sample(num_frames=16, idx=1):\n",
    "    val_dataloader = torch.utils.data.DataLoader(\n",
    "        val_dataset,\n",
    "        batch_size=1,\n",
    "        batch_sampler=batch_sampler,\n",
    "        collate_fn=vis_collate_fn,\n",
    "    )\n",
    "    cnt = 0\n",
    "    for batch in val_dataloader:\n",
    "        cnt += 1\n",
    "        video_pixel_values, video_bboxes, video_cls, video_masks, video_name = batch\n",
    "        all_test_case.append(batch)\n",
    "        width, height = 384, 256\n",
    "        # set gif parameters\n",
    "        gif_name = f'gt_gif_{video_name}_{cnt}.gif'\n",
    "        duration = 500  # 4fps\n",
    "        video_pixel_values, video_bboxes, video_cls, video_masks, video_name = batch\n",
    "        frames = []\n",
    "        color_list = [\"blue\", \"red\", \"blue\", \"green\"]\n",
    "        for fid in range(num_frames):\n",
    "            img_array = ((video_pixel_values.permute(0, 2, 3, 1).numpy()[fid] + 1) * 127.5).astype(np.uint8)\n",
    "            img = Image.fromarray(img_array)\n",
    "\n",
    "            draw = ImageDraw.Draw(img)\n",
    "            _cnt = 0\n",
    "            for bbox in video_bboxes[fid].numpy():\n",
    "                if video_masks[fid][_cnt % 4] == 0:\n",
    "                    _cnt += 1\n",
    "                    continue\n",
    "                x1, y1, x2, y2 = bbox\n",
    "                top_left = (x1 * width, y1 * height)\n",
    "                bottom_right = (x2 * width, y2 * height)\n",
    "                draw.rectangle([top_left, bottom_right], outline=color_list[_cnt % 4], width=3)\n",
    "                _cnt += 1\n",
    "            frames.append(img)\n",
    "\n",
    "        frames[0].save(gif_name, save_all=True, append_images=frames[1:], duration=duration, loop=0, dpi=(300, 300))\n",
    "        break\n",
    "_ = get_fixed_sample()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load model\n",
    "# Enable additional instance embedding\n",
    "os.environ[\"ADD_INS_EMBED\"] = \"true\"\n",
    "# learnable frame-wise time embedding\n",
    "os.environ[\"open_time_embedding\"] = 'learnable_frame'\n",
    "# Disable box attention mechanism\n",
    "os.environ[\"open_box_attn\"] = 'disable'\n",
    "# Disable dynamic size adjustment\n",
    "os.environ[\"dynamic_size\"] = 'disable'\n",
    "# Do not use encoder only\n",
    "os.environ[\"enc_only\"] = 'false'\n",
    "# Do not use decoder only\n",
    "os.environ[\"dec_only\"] = 'false'\n",
    "# Enable injector at decoder fusion stage\n",
    "os.environ[\"injector_on\"] = 'dec_fuse'\n",
    "# Disable self-attention mechanism\n",
    "os.environ[\"self_on\"] = 'false'\n",
    "# Disable track query\n",
    "os.environ[\"track_query\"] = 'false'\n",
    "\n",
    "from diffusers import DPMSolverMultistepScheduler, EulerDiscreteScheduler, DDIMScheduler\n",
    "from diffusers.utils import export_to_video\n",
    "from pipelines.pipeline_text_to_video_synth import TextToVideoSDPipeline\n",
    "from models.unet_3d_condition_gligen import UNet3DConditionModel\n",
    "pretrained_model_path = \"runwayml/stable-diffusion-v1-5\"\n",
    "# Download from here: https://huggingface.co/pengxiang/trackdiffusion_ytvis/tree/main/modelscope_ft/unet\n",
    "unet = UNet3DConditionModel.from_pretrained(\"/path/to/finetuning/unet\", torch_dtype=torch.float16,)\n",
    "\n",
    "pipe = TextToVideoSDPipeline.from_pretrained(pretrained_model_path, unet=unet, torch_dtype=torch.float16, variant=\"fp16\", low_cpu_mem_usage=False)\n",
    "pipe = pipe.to('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('/youtubevis_caption.json', 'r', encoding='utf-8') as file:\n",
    "    caption_data = json.load(file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# video_pixel_values, video_bboxes, video_cls, video_masks = get_random_sample(8, 22)\n",
    "import os\n",
    "from PIL import Image, ImageDraw\n",
    "import numpy as np\n",
    "cnt = 0\n",
    "for batch in all_test_case:\n",
    "    cnt += 1\n",
    "    video_pixel_values, video_bboxes, video_cls, video_masks, video_name = batch\n",
    "    # Flatten the tensor and convert to list of integers\n",
    "    video_cls = video_cls[:16]\n",
    "    video_cls_list = video_cls.view(-1).tolist()\n",
    "\n",
    "    # Convert each integer to its corresponding class\n",
    "    video_cls_str_list = [classes_vis[int_idx] if int_idx < len(classes_vis) else \"\" for int_idx in video_cls_list]\n",
    "    \n",
    "    video_bboxes = video_bboxes[:16]\n",
    "    video_masks = video_masks[:16]\n",
    "    image = pipe(\n",
    "        prompt=caption_data[video_name],\n",
    "        width=384,\n",
    "        height=256,\n",
    "        seg_phrases=video_cls_str_list,\n",
    "        video_masks=video_masks,\n",
    "        bbox_prompt=video_bboxes,\n",
    "        num_frames=16,\n",
    "        num_inference_steps=50,\n",
    "        guidance_scale=1.2, # You may find this strange here, please take a look at our pipeline.\n",
    "    ).frames\n",
    "    \n",
    "    width, height = 384, 256\n",
    "    frames = []\n",
    "    color_list = [\"red\", \"yellow\", \"blue\", \"green\"]\n",
    "    for fid in range(16):\n",
    "        img_array = image[fid]\n",
    "        img = Image.fromarray(img_array)\n",
    "        draw = ImageDraw.Draw(img)\n",
    "        _cnt = 0\n",
    "        for bbox in video_bboxes[fid].numpy()[:]:\n",
    "            x1, y1, x2, y2 = bbox\n",
    "            top_left = (x1 * width, y1 * height)\n",
    "            bottom_right = (x2 * width, y2 * height)\n",
    "            draw.rectangle([top_left, bottom_right], outline=color_list[_cnt % 4], width=2)\n",
    "            _cnt += 1\n",
    "\n",
    "        frames.append(img)\n",
    "\n",
    "    # save as GIF\n",
    "    frames[0].save(f'./output_{video_name}.gif',\n",
    "                format='GIF',\n",
    "                append_images=frames[1:],\n",
    "                save_all=True,\n",
    "                duration=500,\n",
    "                loop=0)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "my_env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
